{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "9d0f0ddf-74bd-4c83-89ff-241ba0c71932",
   "metadata": {},
   "outputs": [],
   "source": [
    "import autogen\n",
    "from autogen.cache import Cache\n",
    "\n",
    "from finrobot.utils import get_current_date, register_keys_from_json\n",
    "from finrobot.data_source import FinnHubUtils, YFinanceUtils\n",
    "from autogen import AssistantAgent, UserProxyAgent\n",
    "\n",
    "from typing import Annotated"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "b26986c5-02c2-4d0c-b660-dd8296a1880b",
   "metadata": {},
   "outputs": [],
   "source": [
    "config_list = [\n",
    "    {\n",
    "        \"model\": \"llama3\",\n",
    "        \"base_url\": \"http://127.0.0.1:4000\",\n",
    "        \"api_key\": \"ollama\",\n",
    "    }\n",
    "]\n",
    "\n",
    "llm_config = {\"config_list\": config_list, \"timeout\": 120, \"temperature\": 0}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "88a0de13-7647-4094-b0e7-e7d85354591a",
   "metadata": {},
   "outputs": [],
   "source": [
    "register_keys_from_json(\"../config_api_keys\")\n",
    "\n",
    "analyst = AssistantAgent(\n",
    "    name=\"Market_Analyst\",\n",
    "    llm_config=llm_config,\n",
    ")\n",
    "\n",
    "user_proxy = UserProxyAgent(\"user_proxy\", code_execution_config=False, max_consecutive_auto_reply=1,  is_termination_msg=lambda x: x.get(\"content\", \"\") and \"TERMINATE\" in x.get(\"content\", \"\"),\n",
    "    human_input_mode=\"NEVER\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "9afb3695-4630-49ca-8bfa-882a48959845",
   "metadata": {},
   "outputs": [],
   "source": [
    "from finrobot.toolkits import register_toolkits\n",
    "tools = [\n",
    "    {\n",
    "        \"function\": YFinanceUtils.get_stock_data,\n",
    "        #\"function\": FinnHubUtils.get_company_news,\n",
    "        \"name\": \"get_stock_news\",\n",
    "        \"description\": \"retrieve stock information related to designated company\"\n",
    "    }\n",
    "]\n",
    "register_toolkits(tools, analyst, user_proxy)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "2c8fac60-c27c-494d-9b23-893099f695b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "company = \"apple\"\n",
    "date='2024-07-01'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "8f191f20-f52c-4cc9-973e-032616e505a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[33muser_proxy\u001b[0m (to Market_Analyst):\n",
      "\n",
      "What is stock price available for apple from 2024-07-01 upon 2024-07-29.\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "[autogen.oai.client: 07-29 20:24:18] {329} WARNING - Model ollama/llama3 is not found. The cost will be 0. In your config_list, add field {\"price\" : [prompt_price_per_1k, completion_token_price_per_1k]} for customized pricing.\n",
      "\u001b[33mMarket_Analyst\u001b[0m (to user_proxy):\n",
      "\n",
      "\u001b[32m***** Suggested tool call (call_c7099725-1961-4d4d-bbc2-fed087087a6f): get_stock_news *****\u001b[0m\n",
      "Arguments: \n",
      "{\"symbol\": \"AAPL\", \"start_date\": \"2024-07-01\", \"end_date\": \"2024-07-29\"}\n",
      "\u001b[32m*******************************************************************************************\u001b[0m\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "\u001b[35m\n",
      ">>>>>>>> EXECUTING FUNCTION get_stock_news...\u001b[0m\n",
      "\u001b[33muser_proxy\u001b[0m (to Market_Analyst):\n",
      "\n",
      "\u001b[33muser_proxy\u001b[0m (to Market_Analyst):\n",
      "\n",
      "\u001b[32m***** Response from calling tool (call_c7099725-1961-4d4d-bbc2-fed087087a6f) *****\u001b[0m\n",
      "                                 Open        High         Low       Close    Volume  Dividends  Stock Splits\n",
      "Date                                                                                                        \n",
      "2024-07-01 00:00:00-04:00  212.089996  217.509995  211.919998  216.750000  60402900        0.0           0.0\n",
      "2024-07-02 00:00:00-04:00  216.149994  220.380005  215.100006  220.270004  58046200        0.0           0.0\n",
      "2024-07-03 00:00:00-04:00  220.000000  221.550003  219.029999  221.550003  37369800        0.0           0.0\n",
      "2024-07-05 00:00:00-04:00  221.649994  226.449997  221.649994  226.339996  60412400        0.0           0.0\n",
      "2024-07-08 00:00:00-04:00  227.089996  227.850006  223.250000  227.820007  59085900        0.0           0.0\n",
      "2024-07-09 00:00:00-04:00  227.929993  229.399994  226.369995  228.679993  48076100        0.0           0.0\n",
      "2024-07-10 00:00:00-04:00  229.300003  233.080002  229.250000  232.979996  62627700        0.0           0.0\n",
      "2024-07-11 00:00:00-04:00  231.389999  232.389999  225.770004  227.570007  64710600        0.0           0.0\n",
      "2024-07-12 00:00:00-04:00  228.919998  232.639999  228.679993  230.539993  53046500        0.0           0.0\n",
      "2024-07-15 00:00:00-04:00  236.479996  237.229996  233.089996  234.399994  62631300        0.0           0.0\n",
      "2024-07-16 00:00:00-04:00  235.000000  236.270004  232.330002  234.820007  43234300        0.0           0.0\n",
      "2024-07-17 00:00:00-04:00  229.449997  231.460007  226.639999  228.880005  57345900        0.0           0.0\n",
      "2024-07-18 00:00:00-04:00  230.279999  230.440002  222.270004  224.179993  66034600        0.0           0.0\n",
      "2024-07-19 00:00:00-04:00  224.820007  226.800003  223.279999  224.309998  49151500        0.0           0.0\n",
      "2024-07-22 00:00:00-04:00  227.009995  227.779999  223.089996  223.960007  48201800        0.0           0.0\n",
      "2024-07-23 00:00:00-04:00  224.369995  226.940002  222.679993  225.009995  39960300        0.0           0.0\n",
      "2024-07-24 00:00:00-04:00  224.000000  224.800003  217.130005  218.539993  61777600        0.0           0.0\n",
      "2024-07-25 00:00:00-04:00  218.929993  220.850006  214.619995  217.490005  51391200        0.0           0.0\n",
      "2024-07-26 00:00:00-04:00  218.699997  219.490005  216.009995  217.960007  41580100        0.0           0.0\n",
      "\u001b[32m**********************************************************************************\u001b[0m\n",
      "\n",
      "--------------------------------------------------------------------------------\n",
      "[autogen.oai.client: 07-29 20:24:28] {329} WARNING - Model ollama/llama3 is not found. The cost will be 0. In your config_list, add field {\"price\" : [prompt_price_per_1k, completion_token_price_per_1k]} for customized pricing.\n",
      "\u001b[33mMarket_Analyst\u001b[0m (to user_proxy):\n",
      "\n",
      "\u001b[32m***** Suggested tool call (call_2ebe01d1-bc44-49a9-9ff7-81d13fa69fe3): get_stock_news *****\u001b[0m\n",
      "Arguments: \n",
      "{\"symbol\": \"AAPL\", \"start_date\": \"2024-07-01\", \"end_date\": \"2024-07-29\"}\n",
      "\u001b[32m*******************************************************************************************\u001b[0m\n",
      "\n",
      "--------------------------------------------------------------------------------\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "ChatResult(chat_id=None, chat_history=[{'content': 'What is stock price available for apple from 2024-07-01 upon 2024-07-29.', 'role': 'assistant'}, {'tool_calls': [{'id': 'call_c7099725-1961-4d4d-bbc2-fed087087a6f', 'function': {'arguments': '{\"symbol\": \"AAPL\", \"start_date\": \"2024-07-01\", \"end_date\": \"2024-07-29\"}', 'name': 'get_stock_news'}, 'type': 'function'}], 'content': None, 'role': 'assistant'}, {'content': '                                 Open        High         Low       Close    Volume  Dividends  Stock Splits\\nDate                                                                                                        \\n2024-07-01 00:00:00-04:00  212.089996  217.509995  211.919998  216.750000  60402900        0.0           0.0\\n2024-07-02 00:00:00-04:00  216.149994  220.380005  215.100006  220.270004  58046200        0.0           0.0\\n2024-07-03 00:00:00-04:00  220.000000  221.550003  219.029999  221.550003  37369800        0.0           0.0\\n2024-07-05 00:00:00-04:00  221.649994  226.449997  221.649994  226.339996  60412400        0.0           0.0\\n2024-07-08 00:00:00-04:00  227.089996  227.850006  223.250000  227.820007  59085900        0.0           0.0\\n2024-07-09 00:00:00-04:00  227.929993  229.399994  226.369995  228.679993  48076100        0.0           0.0\\n2024-07-10 00:00:00-04:00  229.300003  233.080002  229.250000  232.979996  62627700        0.0           0.0\\n2024-07-11 00:00:00-04:00  231.389999  232.389999  225.770004  227.570007  64710600        0.0           0.0\\n2024-07-12 00:00:00-04:00  228.919998  232.639999  228.679993  230.539993  53046500        0.0           0.0\\n2024-07-15 00:00:00-04:00  236.479996  237.229996  233.089996  234.399994  62631300        0.0           0.0\\n2024-07-16 00:00:00-04:00  235.000000  236.270004  232.330002  234.820007  43234300        0.0           0.0\\n2024-07-17 00:00:00-04:00  229.449997  231.460007  226.639999  228.880005  57345900        0.0           0.0\\n2024-07-18 00:00:00-04:00  230.279999  230.440002  222.270004  224.179993  66034600        0.0           0.0\\n2024-07-19 00:00:00-04:00  224.820007  226.800003  223.279999  224.309998  49151500        0.0           0.0\\n2024-07-22 00:00:00-04:00  227.009995  227.779999  223.089996  223.960007  48201800        0.0           0.0\\n2024-07-23 00:00:00-04:00  224.369995  226.940002  222.679993  225.009995  39960300        0.0           0.0\\n2024-07-24 00:00:00-04:00  224.000000  224.800003  217.130005  218.539993  61777600        0.0           0.0\\n2024-07-25 00:00:00-04:00  218.929993  220.850006  214.619995  217.490005  51391200        0.0           0.0\\n2024-07-26 00:00:00-04:00  218.699997  219.490005  216.009995  217.960007  41580100        0.0           0.0', 'tool_responses': [{'tool_call_id': 'call_c7099725-1961-4d4d-bbc2-fed087087a6f', 'role': 'tool', 'content': '                                 Open        High         Low       Close    Volume  Dividends  Stock Splits\\nDate                                                                                                        \\n2024-07-01 00:00:00-04:00  212.089996  217.509995  211.919998  216.750000  60402900        0.0           0.0\\n2024-07-02 00:00:00-04:00  216.149994  220.380005  215.100006  220.270004  58046200        0.0           0.0\\n2024-07-03 00:00:00-04:00  220.000000  221.550003  219.029999  221.550003  37369800        0.0           0.0\\n2024-07-05 00:00:00-04:00  221.649994  226.449997  221.649994  226.339996  60412400        0.0           0.0\\n2024-07-08 00:00:00-04:00  227.089996  227.850006  223.250000  227.820007  59085900        0.0           0.0\\n2024-07-09 00:00:00-04:00  227.929993  229.399994  226.369995  228.679993  48076100        0.0           0.0\\n2024-07-10 00:00:00-04:00  229.300003  233.080002  229.250000  232.979996  62627700        0.0           0.0\\n2024-07-11 00:00:00-04:00  231.389999  232.389999  225.770004  227.570007  64710600        0.0           0.0\\n2024-07-12 00:00:00-04:00  228.919998  232.639999  228.679993  230.539993  53046500        0.0           0.0\\n2024-07-15 00:00:00-04:00  236.479996  237.229996  233.089996  234.399994  62631300        0.0           0.0\\n2024-07-16 00:00:00-04:00  235.000000  236.270004  232.330002  234.820007  43234300        0.0           0.0\\n2024-07-17 00:00:00-04:00  229.449997  231.460007  226.639999  228.880005  57345900        0.0           0.0\\n2024-07-18 00:00:00-04:00  230.279999  230.440002  222.270004  224.179993  66034600        0.0           0.0\\n2024-07-19 00:00:00-04:00  224.820007  226.800003  223.279999  224.309998  49151500        0.0           0.0\\n2024-07-22 00:00:00-04:00  227.009995  227.779999  223.089996  223.960007  48201800        0.0           0.0\\n2024-07-23 00:00:00-04:00  224.369995  226.940002  222.679993  225.009995  39960300        0.0           0.0\\n2024-07-24 00:00:00-04:00  224.000000  224.800003  217.130005  218.539993  61777600        0.0           0.0\\n2024-07-25 00:00:00-04:00  218.929993  220.850006  214.619995  217.490005  51391200        0.0           0.0\\n2024-07-26 00:00:00-04:00  218.699997  219.490005  216.009995  217.960007  41580100        0.0           0.0'}], 'role': 'tool'}, {'tool_calls': [{'id': 'call_2ebe01d1-bc44-49a9-9ff7-81d13fa69fe3', 'function': {'arguments': '{\"symbol\": \"AAPL\", \"start_date\": \"2024-07-01\", \"end_date\": \"2024-07-29\"}', 'name': 'get_stock_news'}, 'type': 'function'}], 'content': None, 'role': 'assistant'}], summary='', cost={'usage_including_cached_inference': {'total_cost': 0, 'ollama/llama3': {'cost': 0, 'prompt_tokens': 1375, 'completion_tokens': 97, 'total_tokens': 1472}}, 'usage_excluding_cached_inference': {'total_cost': 0, 'ollama/llama3': {'cost': 0, 'prompt_tokens': 1375, 'completion_tokens': 97, 'total_tokens': 1472}}}, human_input=[])"
      ]
     },
     "execution_count": 18,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "user_proxy.initiate_chat(analyst, message=f\"What is stock price available for {company} from {date} upon {get_current_date()}.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7dd350c5-bd56-48ae-875d-7c6d2d146f7a",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
